#include "resnet50.h"
//#include "utils/Logger.h"
#include <algorithm>

Resnet50Inference::Resnet50Inference() {
	//data = nullptr;
	//prob = nullptr;
	hasCreate = false;

	context = nullptr;
	engine = nullptr;
	runtime = nullptr;
}

Resnet50Inference::~Resnet50Inference() {
	//if (data != nullptr) {
	//	delete[] data;
	//	data = nullptr;
	//}
	//if (prob != nullptr) {
	//	delete[] prob;
	//	prob = nullptr;
	//}
	// Release stream and buffers
	if (hasCreate) {
		cudaStreamDestroy(stream);
		CUDA_CHECK(cudaFree(buffers[inputIndex]));
		CUDA_CHECK(cudaFree(buffers[outputIndex]));
	}

	// Destroy the engine
	if (context != nullptr) {
		context->destroy();
		context = nullptr;
	}
	if (engine != nullptr) {
		engine->destroy();
		engine = nullptr;
	}
	if (runtime != nullptr) {
		runtime->destroy();
		runtime = nullptr;
	}
}

void Resnet50Inference::init() {
#define USE_FP16
	gpu_device_ = 0;
	batch_size_ = 50;
	input_img_w_ = 224;
	input_img_h_ = 224;
	class_num_ = 2;
	output_size_ = class_num_;

	input_blob_name_ = std::string("data");//？？
	output_blob_name_ = std::string("prob");//？？
}

bool Resnet50Inference::loadModel(std::string engine_name) {
	std::ifstream file(engine_name, std::ios::binary);
	if (!file.good()) {
		std::cerr << "read " << engine_name << " error!" << std::endl;
		return false;
	}
	char *trtModelStream = nullptr;
	file.seekg(0, file.end);
	size_t size = file.tellg();
	file.seekg(0, file.beg);
	trtModelStream = new char[size];
	assert(trtModelStream);
	file.read(trtModelStream, size);
	file.close();


	runtime = nvinfer1::createInferRuntime(gLogger);
	assert(runtime != nullptr);
	engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
	assert(engine != nullptr);
	context = engine->createExecutionContext();
	assert(context != nullptr);
	delete[] trtModelStream;
	assert(engine->getNbBindings() == 2);

	inputIndex = engine->getBindingIndex(input_blob_name_.c_str());
	outputIndex = engine->getBindingIndex(output_blob_name_.c_str());
	assert(inputIndex == 0);
	assert(outputIndex == 1);
	// Create GPU buffers on device
	CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batch_size_ * 3 * input_img_h_ * input_img_w_ * sizeof(float)));
	CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batch_size_ * output_size_ * sizeof(float)));
	// Create stream
	CUDA_CHECK(cudaStreamCreate(&stream));
	hasCreate = true;
	return true;
}

void Resnet50Inference::set_gpu_device_(int _gpu_device_) {
	gpu_device_ = _gpu_device_;
}

void Resnet50Inference::set_batch_size(int _batch_size_) {
	batch_size_ = _batch_size_;
}

void Resnet50Inference::set_input_img_w(int _input_img_w_) {
	input_img_w_ = _input_img_w_;
}

void Resnet50Inference::set_input_img_h_(int _input_img_h_) {
	input_img_h_ = _input_img_h_;
}

/****************empty****************************/
void Resnet50Inference::set_conf_thresh(float _conf_thresh_) {
}

std::vector<YOLOResult> Resnet50Inference:: detect(cv::Mat img) {
	std::vector<YOLOResult>a;
	return a;
}

void Resnet50Inference::set_class_num(int _class_num_) {
	class_num_ = _class_num_;
	output_size_ = _class_num_;
}

void Resnet50Inference::addCategories(std::string category) {
	categories.push_back(category);
}

void Resnet50Inference::doInference(nvinfer1::IExecutionContext& context, float* input, float* output, int batchSize) {
	CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * input_img_h_ * input_img_w_ * sizeof(float), cudaMemcpyHostToDevice, stream));
	context.enqueue(batchSize, buffers, stream, nullptr);
	CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * output_size_ * sizeof(float), cudaMemcpyDeviceToHost, stream));
	cudaStreamSynchronize(stream);
}

cv::Mat Resnet50Inference::preprocess_img(cv::Mat& img, int input_w, int input_h) {
	int w, h, x, y;
	int shape_x = img.cols / 2;
	int shape_y = img.rows / 2;
	x = shape_x - input_w / 2;
	y = shape_y - input_h / 2;
	cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128));
	img(cv::Rect(x, y, input_w, input_h)).copyTo(out);
	return out;
}

//int Resnet50Inference::classify(cv::Mat img) {
//	//对输入图进行resize，拷贝到img640，最后resize到input_img_w_、input_img_h_尺寸
//	cv::Mat img256 = cv::Mat(256, 256, CV_8UC3, cv::Scalar(255, 255, 255));
//	cv::Point2d centerPoint(128, 128);
//	int startx = centerPoint.x - img.cols / 2;
//	int starty = centerPoint.y - img.rows / 2;
//	int endx = centerPoint.x + img.cols / 2;
//	int endy = centerPoint.y + img.rows / 2;
//	startx = startx > 0 ? startx : 0;
//	starty = starty > 0 ? starty : 0;
//	if (endx > 256 || endy > 256) {
//		endx = endx > 256 ? 256 : endx;
//		endy = endy > 256 ? 256 : endy;
//		int width = endx - startx;
//		int height = endy - starty;
//		cv::resize(img, img, cv::Size(width, height));
//	}
//	cv::Rect objRect(startx, starty, img.cols, img.rows);
//	img.copyTo(img256(objRect));
//	cv::resize(img256, img, cv::Size(input_img_w_, input_img_h_), 0, 0, cv::INTER_LINEAR);
//
//	//cv::resize(img, img, cv::Size(input_img_w_, input_img_h_), cv::INTER_LINEAR);
//	std::vector<float> mean_value{ 0.485f, 0.456f, 0.406f };//这组数值如何确定？voc数据集？
//	std::vector<float> std_value{ 0.229f, 0.224f, 0.225f };
//	cv::Mat dst;
//	std::vector<cv::Mat> rgbChannels(3);
//	cv::split(img, rgbChannels);
//
//	for (auto i = 0; i < rgbChannels.size(); i++){
//		rgbChannels[i].convertTo(rgbChannels[i], CV_32FC1, 1.0 / (std_value[i] * 255.0), (0.0 - mean_value[i]) / std_value[i]);
//	}
//	cv::merge(rgbChannels, dst);
//
//	float* data = new float[input_img_h_*input_img_w_ * 3];
//	float* prob = new float[class_num_];
//
//	int i = 0;
//	float val;
//	
//	for (int row = 0; row < input_img_h_; ++row) {
//		for (int col = 0; col < input_img_w_; ++col) {
//			data[i] = dst.at<cv::Vec3f>(row, col)[0];
//			data[i + input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[1];
//			data[i + 2 * input_img_h_ * input_img_h_] = dst.at<cv::Vec3f>(row, col)[2];
//			++i;
//		}
//	}
//	doInference(*context, data, prob, 1);
//	int ret = 0;
//	if (prob[0] > prob[1]) {
//		ret = 0;
//	}
//	else {
//		ret = 1;
//	}
//
//	delete data;
//	delete prob;
//	data = nullptr;
//	prob = nullptr;
//	return ret;
//}

int Resnet50Inference::classify(cv::Mat img) {
	//cv::resize(img, img, cv::Size(input_img_w_, input_img_h_), 0, 0, cv::INTER_LINEAR);

	cv::resize(img, img, cv::Size(input_img_w_, input_img_h_), cv::INTER_LINEAR);
	std::vector<float> mean_value{ 0.485f, 0.456f, 0.406f };
	std::vector<float> std_value{ 0.229f, 0.224f, 0.225f };
	cv::Mat dst;
	std::vector<cv::Mat> rgbChannels(3);
	cv::split(img, rgbChannels);

	for (auto i = 0; i < rgbChannels.size(); i++)
	{
		rgbChannels[i].convertTo(rgbChannels[i], CV_32FC1, 1.0 / (std_value[i] * 255.0), (0.0 - mean_value[i]) / std_value[i]);
	}
	cv::merge(rgbChannels, dst);

	float* data = new float[input_img_h_*input_img_w_ * 3];
	float* prob = new float[class_num_];

	/*int i = 0;
	float val;
	for (int row = 0; row < input_img_h_; ++row) {
		float* uc_pixel = dst.ptr<float>(row);
		for (int col = 0; col < input_img_w_; ++col) {
			data[i] = uc_pixel[0];
			data[i + input_img_h_ * input_img_w_] = uc_pixel[1];
			data[i + 2 * input_img_h_ * input_img_h_] = uc_pixel[2];
			++i;
			uc_pixel += 3;
		}
	}*/

	int i = 0;
	float val;
	for (int row = 0; row < input_img_h_; ++row) {
		float* uc_pixel = dst.ptr<float>(row);
		for (int col = 0; col < input_img_w_; ++col) {
			data[i] = uc_pixel[0];
			data[i + input_img_h_ * input_img_w_] = uc_pixel[1];
			data[i + 2 * input_img_h_ * input_img_h_] = uc_pixel[2];
			++i;
			uc_pixel += 3;
		}
	}


	//int i = 0;
	//float val;
	//for (int row = 0; row < input_img_h_; ++row) {
	//	//uchar* uc_pixel = dst.ptr<float>data + row * pr_img.step;
	//	for (int col = 0; col < input_img_w_; ++col) {
	//		data[i] = dst.at<cv::Vec3f>(row, col)[0];
	//		data[i + input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[1];
	//		data[i + 2 * input_img_h_ * input_img_h_] = dst.at<cv::Vec3f>(row, col)[2];
	//		++i;
	//	}
	//}
	doInference(*context, data, prob, 1);
	int ret = 0;
	ret = std::max_element(prob, prob + class_num_) - prob;

	delete data;
	delete prob;
	data = nullptr;
	prob = nullptr;
	return ret;
}

//std::vector<int> Resnet50Inference::classify(std::vector<cv::Mat> imgs) {
//	std::vector<int> classifyIdxs;
//	int startIdx = 0;
//	while (startIdx <imgs.size())
//	{
//		int cur_batch_size = batch_size_;
//		if ((startIdx + batch_size_) > imgs.size()) {
//			cur_batch_size = imgs.size() - startIdx;
//			//startIdx += cur_batch_size;
//		}
//
//		float* data = new float[cur_batch_size * input_img_h_*input_img_w_ * 3];
//		float* prob = new float[cur_batch_size * class_num_];
//
//		for (size_t idx = 0; idx < cur_batch_size; idx++)
//		{
//			cv::Mat img = imgs[idx + startIdx];
//			cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
//
//			cv::Mat pr_img;
//			cv::resize(img, img, cv::Size(256, 256), cv::INTER_LINEAR);
//			//center crop
//			pr_img = preprocess_img(img, input_img_w_, input_img_h_); // letterbox BGR to RGB
//			pr_img.convertTo(pr_img, CV_32FC3);
//			pr_img = pr_img / 255.0;
//
//			std::vector<float> mean_value{ 0.485f, 0.456f, 0.406f };
//			std::vector<float> std_value{ 0.229f, 0.224f, 0.225f };
//			cv::Mat dst;
//			std::vector<cv::Mat> rgbChannels(3);
//			cv::split(pr_img, rgbChannels);
//			for (auto i = 0; i < rgbChannels.size(); i++)
//			{
//				rgbChannels[i] = (rgbChannels[i] - mean_value[i]) / std_value[i];
//			}
//			cv::merge(rgbChannels, dst);
//			int i = 0;
//			float val;
//			for (int row = 0; row < input_img_h_; ++row) {
//				for (int col = 0; col < input_img_w_; ++col) {
//					data[idx * 3 * input_img_h_ * input_img_w_ + i] = dst.at<cv::Vec3f>(row, col)[0];
//					data[idx * 3 * input_img_h_ * input_img_w_  + i + input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[1];
//					data[idx * 3 * input_img_h_ * input_img_w_ + i + 2 * input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[2];
//					++i;
//				}
//			}
//		}
//		doInference(*context, data, prob, cur_batch_size);
//		//LInfo("before fill  classifyIdxs detect");
//
//		for (size_t i = 0; i < cur_batch_size; i++)
//		{
//			float* cur_prob = prob + i * output_size_;
//			auto maxPosition = std::max_element(cur_prob, cur_prob + output_size_) - cur_prob;
//			classifyIdxs.push_back(maxPosition);
//		}
//		//LInfo("after fill  classifyIdxs detect");
//
//		delete data;
//		delete prob;
//		data = nullptr;
//		prob = nullptr;
//
//		startIdx += cur_batch_size;
//	}
//	return classifyIdxs;
//}

std::vector<int> Resnet50Inference::classify(std::vector<cv::Mat> imgs) {
	std::vector<int> classifyIdxs;
	int startIdx = 0;
	while (startIdx <imgs.size())
	{
		int cur_batch_size = batch_size_;
		if ((startIdx + batch_size_) > imgs.size()) {
			cur_batch_size = imgs.size() - startIdx;
			//startIdx += cur_batch_size;
		}

		float* data = new float[cur_batch_size * input_img_h_*input_img_w_ * 3];
		float* prob = new float[cur_batch_size * class_num_];

		for (size_t idx = 0; idx < cur_batch_size; idx++)
		{
			cv::Mat img = imgs[idx + startIdx];
			cv::cvtColor(img, img, cv::COLOR_BGR2RGB);

			cv::Mat pr_img;
			cv::resize(img, img, cv::Size(224, 224), cv::INTER_LINEAR);
			//center crop
			pr_img = preprocess_img(img, input_img_w_, input_img_h_); // letterbox BGR to RGB
			pr_img.convertTo(pr_img, CV_32FC3);
			pr_img = pr_img / 255.0;

			std::vector<float> mean_value{ 0.485f, 0.456f, 0.406f };
			std::vector<float> std_value{ 0.229f, 0.224f, 0.225f };
			cv::Mat dst;
			std::vector<cv::Mat> rgbChannels(3);
			cv::split(pr_img, rgbChannels);
			for (auto i = 0; i < rgbChannels.size(); i++)
			{
				rgbChannels[i] = (rgbChannels[i] - mean_value[i]) / std_value[i];
			}
			cv::merge(rgbChannels, dst);
			//int i = 0;
			//float val;
			//for (int row = 0; row < input_img_h_; ++row) {
			//	for (int col = 0; col < input_img_w_; ++col) {
			//		data[idx * 3 * input_img_h_ * input_img_w_ + i] = dst.at<cv::Vec3f>(row, col)[0];
			//		data[idx * 3 * input_img_h_ * input_img_w_  + i + input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[1];
			//		data[idx * 3 * input_img_h_ * input_img_w_ + i + 2 * input_img_h_ * input_img_w_] = dst.at<cv::Vec3f>(row, col)[2];
			//		++i;
			//	}
			//}	
			int i = 0;
			float val;
			for (int row = 0; row < input_img_h_; ++row) {
				float* uc_pixel = dst.ptr<float>(row);
				for (int col = 0; col < input_img_w_; ++col) {
					data[idx * 3 * input_img_h_ * input_img_w_ + i] = uc_pixel[0];
					data[idx * 3 * input_img_h_ * input_img_w_ + i + input_img_h_ * input_img_w_] = uc_pixel[1];
					data[idx * 3 * input_img_h_ * input_img_w_ + i + 2 * input_img_h_ * input_img_h_] = uc_pixel[2];
					++i;
					uc_pixel += 3;
				}
			}
		}
		doInference(*context, data, prob, cur_batch_size);
		//LInfo("before fill  classifyIdxs detect");

		for (size_t i = 0; i < cur_batch_size; i++)
		{
			float* cur_prob = prob + i * output_size_;
			auto maxPosition = std::max_element(cur_prob, cur_prob + output_size_) - cur_prob;
			classifyIdxs.push_back(maxPosition);
		}
		//LInfo("after fill  classifyIdxs detect");

		delete data;
		delete prob;
		data = nullptr;
		prob = nullptr;

		startIdx += cur_batch_size;
	}
	return classifyIdxs;
}

